from sklearn.metrics import confusion_matrix
import seaborn as sns
import torch
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from Config import config  # 从自定义的Config.py参数文件中插入
from Model import Model

# 加载参数
config = config()

# ToTenser():1、转化为一个tensor,2、将数值归一化,转换到9-1之间3、将channel维度放在第一维度上
transformation = transforms.Compose([
    transforms.ToTensor()
])

test_ds = datasets.MNIST(config.data_path, train=False, download=True, transform=transformation)
test_dl = DataLoader(test_ds, batch_size=config.batch_size, shuffle=True)
# print(len(test_ds))
# 载入模型和参数
model = Model()  # 初始化模型
model.load_state_dict(torch.load(config.save_path))  # 导入config 类中的参数
model = model.eval()  # 消除dropout 层的影响

predicted_labels = []
true_labels = []
# 将图片输入到模型中输出预测矩阵
with torch.no_grad():
    for images, labels in test_dl:
        # print(images, shape)
        outputs = model(images)

        # torch.argmax(y_pred,dim=1):找到每个样本图片输出概率最高的类别,dim=1表示在类别维度上进行操作
        # .numpy()[i]:将张量转换为numpy数组,并提取第i个样本的预测类别
        predicted_labels.extend(torch.argmax(outputs, dim=1).numpy())
        true_labels.extend(labels.numpy())
#计算混淆矩阵
cm = confusion_matrix(true_labels, predicted_labels)

#使用Seaborn绘制混淆矩阵热图
plt.figure(figsize=(8, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=np.arange(10), yticklabels=np.arange(10))
plt.xlabel('Predicted')
plt.ylabel('True')
plt.title('Confusion Matrix')
plt.show()